package cn.kennylee.codehub.springboot.microtrace.thread;

import cn.hutool.json.JSONUtil;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.core.Ordered;
import org.springframework.core.annotation.Order;
import org.springframework.core.task.TaskDecorator;
import org.springframework.lang.NonNull;

import java.util.Map;

/**
 * 异步线程池的MDC支持
 * <p>Created on 2025/4/30.</p>
 *
 * @author kennylee
 * @since 0.0.1
 */
@Slf4j
@Order(Ordered.HIGHEST_PRECEDENCE)
public class MdcTaskDecorator implements TaskDecorator {

    @Override
    @NonNull
    public Runnable decorate(@NonNull Runnable runnable) {
        // 保存当前线程的 MDC 上下文
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        if (contextMap == null || contextMap.isEmpty()) {
            return runnable;
        }

        return () -> {
            try {
                if (log.isDebugEnabled()) {
                    log.debug("设置 MDC 上下文: {}", JSONUtil.toJsonStr(contextMap));
                }
                MDC.setContextMap(contextMap);
                runnable.run();
            } finally {
                MDC.clear();
            }
        };
    }
}
